# models/CS_KGAT.py
import torch
import torch.nn as nn
import torch.nn.functional as F


def _L2_loss_mean(x):
    """计算L2正则化损失的辅助函数。"""
    return torch.mean(torch.sum(torch.pow(x, 2), dim=1, keepdim=False) / 2.)


class ConfidenceAwareAggregator(nn.Module):
    """
    置信度感知的图聚合层。
    此聚合器借鉴了 KGAT 的 Bi-Interaction 设计，通过在消息传播中直接使用
    一个已经包含了LLM置信度作为权重的邻接矩阵 A_in，从而实现了对知识可靠性
    的感知。torch.sparse.mm 运算会自动利用这些权重来加权邻居节点的消息。
    """

    def __init__(self, in_dim, out_dim, dropout):
        super().__init__()
        self.message_dropout = nn.Dropout(dropout)
        # LeakyReLU 是GNN中常用的激活函数
        self.activation = nn.LeakyReLU(negative_slope=0.01)

        # Bi-Interaction 的两个线性变换层
        self.linear1 = nn.Linear(in_dim, out_dim)
        self.linear2 = nn.Linear(in_dim, out_dim)

        # 使用 Xavier 初始化权重
        nn.init.xavier_uniform_(self.linear1.weight)
        nn.init.xavier_uniform_(self.linear2.weight)

    def forward(self, ego_embeddings, A_in):
        """
        前向传播。
        :param ego_embeddings: (n_nodes, in_dim) - 所有节点自身的嵌入。
        :param A_in: (n_nodes, n_nodes) - 带权重的稀疏邻接矩阵。
        """
        # 消息传播。由于A_in是加权且归一化的，此步骤已隐式包含置信度。
        side_embeddings = torch.sparse.mm(A_in, ego_embeddings)

        # Bi-Interaction 聚合方式 (来自 KGAT)
        # 1. 元素相加变换
        sum_embeddings = self.linear1(ego_embeddings + side_embeddings)
        # 2. 元素相乘变换
        bi_embeddings = self.linear2(ego_embeddings * side_embeddings)

        # 聚合两种交互信息
        embeddings = self.activation(sum_embeddings + bi_embeddings)

        return self.message_dropout(embeddings)


class CS_KGAT(nn.Module):
    """
    CSE-KGAN的核心模型: 置信度感知与语义增强的知识图谱注意力网络。
    该模型深度融合了三个核心思想:
    1.  **KGAT**: 使用GNN学习知识图谱中的高阶结构化信息。
    2.  **EasyRec**: 引入LLM生成的语义画像来丰富节点表示。
    3.  **CSRec/CKG-LLMA**: 在GNN中融入LLM生成的常识关系及其置信度，并通过对比学习对齐多视图信息。
    """

    # 【代码优化】: 直接传入data对象，使参数传递更简洁
    def __init__(self, args, data, user_profiles_embed, item_profiles_embed):
        super().__init__()

        # --- 1. 初始化基本参数 ---
        self.args = args
        self.n_users = data.n_users
        self.n_items = data.n_items
        self.n_entities = data.n_entities
        self.n_nodes = data.n_nodes
        self.device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() and args.gpu_id != -1 else "cpu")

        # --- 2. 结构化信息学习模块 (GNN部分) ---
        self.embed_dim = args.embed_dim
        self.conv_dim_list = [args.embed_dim] + eval(args.conv_dim_list)
        self.n_layers = len(self.conv_dim_list) - 1
        self.mess_dropout = eval(args.mess_dropout)

        # 节点嵌入（包括用户和所有实体）
        self.all_embed = nn.Embedding(self.n_nodes, self.embed_dim)
        nn.init.xavier_uniform_(self.all_embed.weight)

        # GNN 聚合层列表
        self.aggregator_layers = nn.ModuleList()
        for k in range(self.n_layers):
            self.aggregator_layers.append(
                ConfidenceAwareAggregator(self.conv_dim_list[k], self.conv_dim_list[k + 1], self.mess_dropout[k])
            )

        # 从DataLoader获取的、已包含置信度权重的邻接矩阵
        self.A_in = data.A_in.to(self.device)

        # --- 3. 语义信息增强模块 (LLM部分) ---
        # 加载预先通过LLM生成并编码好的画像嵌入
        # freeze=True 表示这些嵌入在训练中是固定的，不参与反向传播，作为一种先验知识
        self.user_profiles_embed = nn.Embedding.from_pretrained(user_profiles_embed.float(), freeze=True).to(
            self.device)
        self.item_profiles_embed = nn.Embedding.from_pretrained(item_profiles_embed.float(), freeze=True).to(
            self.device)

        # 【核心创新】语义投影层
        # 用于将LLM画像的语义嵌入映射到与GNN最后一层相同的维度空间，以便进行有效融合。
        semantic_embed_dim = self.user_profiles_embed.embedding_dim
        # 【鲁棒性优化】: 结构化嵌入的总维度
        structure_total_dim = sum(self.conv_dim_list)

        self.semantic_projection = nn.Linear(semantic_embed_dim, structure_total_dim)
        nn.init.xavier_uniform_(self.semantic_projection.weight)

        # --- 4. 损失函数超参数 ---
        self.cf_l2loss_lambda = args.cf_l2loss_lambda
        self.contrastive_lambda = args.contrastive_lambda

    def _get_gnn_embeddings(self):
        """通过GNN层计算所有节点的结构化嵌入。"""
        ego_embed = self.all_embed.weight
        all_embeds = [ego_embed]  # 存储每一层的嵌入，用于最终拼接

        for aggregator in self.aggregator_layers:
            ego_embed = aggregator(ego_embed, self.A_in)
            all_embeds.append(ego_embed)

        # 类似于KGAT，拼接所有层的嵌入作为最终的结构化表示
        # 维度: (n_nodes, embed_dim + conv_dim_1 + ... + conv_dim_L)
        return torch.cat(all_embeds, dim=1)

    def calculate_loss(self, users, pos_items, neg_items):
        """
        计算模型的总损失，集成了BPR损失、正则化损失和对比损失。
        :param users: (batch_size) 用户ID
        :param pos_items: (batch_size) 正样本物品ID
        :param neg_items: (batch_size) 负样本物品ID
        """
        # 在计算图中，用户ID需要加上实体数量的偏移
        user_g_ids = users + self.n_entities

        # --- 1. 获取结构化嵌入 (从GNN) ---
        all_structural_embeds = self._get_gnn_embeddings()
        user_struct_embed = all_structural_embeds[user_g_ids]
        pos_item_struct_embed = all_structural_embeds[pos_items]
        neg_item_struct_embed = all_structural_embeds[neg_items]

        # --- 2. 获取语义嵌入 (从LLM画像) ---
        # 注意：这里的ID是原始ID，可以直接用于索引画像嵌入
        user_sem_embed = self.user_profiles_embed(users)
        pos_item_sem_embed = self.item_profiles_embed(pos_items)
        neg_item_sem_embed = self.item_profiles_embed(neg_items)

        # 将语义嵌入投影到与GNN结构嵌入相同的维度空间
        user_sem_proj = self.semantic_projection(user_sem_embed)
        pos_item_sem_proj = self.semantic_projection(pos_item_sem_embed)
        neg_item_sem_proj = self.semantic_projection(neg_item_sem_embed)

        # --- 3. 【核心创新】融合结构与语义嵌入 ---
        # 简单的加法融合。
        user_final_embed = user_struct_embed + user_sem_proj
        pos_item_final_embed = pos_item_struct_embed + pos_item_sem_proj
        neg_item_final_embed = neg_item_struct_embed + neg_item_sem_proj

        # --- 4. 损失计算 ---
        # a) BPR推荐损失 (核心任务)
        pos_scores = torch.sum(user_final_embed * pos_item_final_embed, dim=1)
        neg_scores = torch.sum(user_final_embed * neg_item_final_embed, dim=1)
        cf_loss = -torch.mean(F.logsigmoid(pos_scores - neg_scores))

        # b) L2正则化损失
        l2_loss = _L2_loss_mean(user_final_embed) + _L2_loss_mean(pos_item_final_embed) + \
                  _L2_loss_mean(neg_item_final_embed)

        # c) 【核心创新】对比损失 (InfoNCE)
        # 目标是让同一个节点的结构表示和语义表示在向量空间中更接近，从而实现视图对齐。
        cl_loss_user = self._calculate_contrastive_loss(user_struct_embed, user_sem_proj)
        cl_loss_item = self._calculate_contrastive_loss(pos_item_struct_embed, pos_item_sem_proj)
        cl_loss = cl_loss_user + cl_loss_item

        # d) 总损失
        total_loss = cf_loss + self.cf_l2loss_lambda * l2_loss + self.contrastive_lambda * cl_loss
        return total_loss, cf_loss, l2_loss, cl_loss

    def _calculate_contrastive_loss(self, view1, view2, temperature=0.1):
        """计算InfoNCE对比损失。"""
        # L2归一化
        view1 = F.normalize(view1, dim=1)
        view2 = F.normalize(view2, dim=1)

        # 计算相似度矩阵 (batch_size, batch_size)
        sim_matrix = torch.matmul(view1, view2.T) / temperature

        # 对角线上的元素是正样本对，其余为负样本对
        labels = torch.arange(view1.size(0)).to(self.device)

        return F.cross_entropy(sim_matrix, labels)

    def predict(self, users, items):
        """
        用于评估的前向传播，计算用户对一批物品的预测分数。
        :param users: (n_eval_users) 待评估的用户ID
        :param items: (n_items) 所有物品ID, 此处为None，表示对所有物品进行预测
        """
        # 在评估时，我们使用完整的图，不进行dropout
        with torch.no_grad():
            all_structural_embeds = self._get_gnn_embeddings()

            # 获取所有用户和物品的语义嵌入并投影
            all_user_sem_embed = self.user_profiles_embed.weight
            all_item_sem_embed = self.item_profiles_embed.weight
            all_user_sem_proj = self.semantic_projection(all_user_sem_embed)
            all_item_sem_proj = self.semantic_projection(all_item_sem_embed)

            # 提取结构嵌入
            user_struct_embeds = all_structural_embeds[self.n_entities:]
            item_struct_embeds = all_structural_embeds[:self.n_items]

            # 融合
            user_final_embeds = user_struct_embeds + all_user_sem_proj
            item_final_embeds = item_struct_embeds + all_item_sem_proj

        # --- 2. 提取当前批次用户的嵌入 ---
        batch_user_embeds = user_final_embeds[users]

        # --- 3. 计算预测分数 (点积) ---
        scores = torch.matmul(batch_user_embeds, item_final_embeds.T)
        return scores